import random
import numpy as np
import networkx as nx


def random_pcn(node_num, density):
    g = np.full([node_num, node_num], 0)

    for i in range(node_num):
        for j in range(i + 1, node_num):
            if random.randint(0, node_num) < density:
                g[i][j] = 1
                g[j][i] = 1

    return g


def ws_pcn(node_num, neighbor, poss):
    ws = nx.watts_strogatz_graph(node_num, neighbor, poss)
    pcn = nx.to_numpy_array(ws)
    return pcn


def pcn_path(pcn, steps):
    l = len(pcn)
    path = pcn.copy()
    now = pcn.copy()
    for i in range(2, steps+1):
        now = np.matmul(now, pcn)

        for j in range(l):
            for k in range(l):
                if path[j][k] == 0 and now[j][k] != 0:
                    path[j][k] = i # 记录该路径的长度

    return path


def get_node_path(node_num, graph, steps, now):
    path = [[] for _ in range(node_num)]
    one = []
    two = []
    three = []

    # reach in one step
    for i in range(node_num):
        if graph[now][i] != 0:
            path[i] = [i]
            one.append(i)

    # reach in two step
    if steps >= 2:
        for i in range(len(one)):
            for j in range(node_num):
                if len(path[j]) != 0 and graph[one[i]][j] != 0:
                    path[j] = [one[i], j]
                    two.append(j)

    # reach in three step
    if steps >= 3:
        for i in range(len(two)):
            for j in range(node_num):
                if len(path[j]) != 0 and graph[two[i]][j] != 0:
                    path[j] = path[two[i]] + [j]
                    three.append(j)

    return path


def find_path(node_num, graph, steps):
    path = []

    for i in range(node_num):
        now = i
        path.append(get_node_path(node_num, graph, steps, now))

    return path


# print(ws_pcn(10, 3, 0.5))